import pandas as pd
import os
import json
from collections import defaultdict
"""
    generate pid csv
"""
input_file = '/home/chenjiaxiong/jx_code/a3ddd04fd63055e1529e2f8782db498e/a3ddd04fd63055e1529e2f8782db498e.csv'  # 替换为你的CSV文件名
data = pd.read_csv(input_file)

# # 检查是否存在输出目录，没有就创建
output_dir = '/home/chenjiaxiong/jx_code/a3ddd04fd63055e1529e2f8782db498e'
os.makedirs(output_dir, exist_ok=True)

# 遍历每个唯一的PID，并生成对应的CSV文件
for pid in data['PID'].unique():
    # 过滤出该PID对应的数据
    pid_data = data[data['PID'] == pid]
    
    # 生成文件名
    output_file = f'{output_dir}/PID_{pid}.csv'
    
    # 保存为单独的CSV文件
    pid_data.to_csv(output_file, index=False)

print(f"所有PID的CSV文件已生成并保存在 '{output_dir}' 文件夹中。")

"""
    generate process relation
"""

process_creation_calls = ['clone', 'fork', 'vfork', 'posix_spawn']
# 创建一个字典来存储父子关系
process_tree = defaultdict(list)

# 遍历所有 'clone' 行，构建嵌套结构
for _, row in data.iterrows():
    if row['Function'] in process_creation_calls and pd.notna(row['Result']):
        parent_pid = row['PID']
        child_pid = int(row['Result'].split()[0])
        process_tree[parent_pid].append(child_pid)

# 将字典转换为 JSON 并保存
output_file = f'{output_dir}/process_tree.json'
with open(output_file, 'w') as f:
    json.dump(process_tree, f, indent=4)

print(f"进程树已生成并保存为 '{output_file}'。")

"""
    generate process graph
"""
import pandas as pd
import re
import networkx as nx
import matplotlib.pyplot as plt

# 读取CSV文件
df = pd.read_csv("/home/chenjiaxiong/jx_code/a3ddd04fd63055e1529e2f8782db498e/PID_31755.csv")  # 替换为你的CSV文件路径

# 使用改进后的正则表达式匹配文件描述符
fd_pattern = re.compile(r"(\d+<.*>)")
param_fds_pattern = re.compile(r"(\d+<[^<>]*?>)")

# 初始化有向图和 fd 列表
G = nx.DiGraph()
fd_dict = {} 

# 遍历每一行记录
for index, row in df.iterrows():
    # 尝试从 Result 中提取完整的 fd
    result_fd = fd_pattern.search(str(row['Result']))

    if result_fd:  # 如果找到了有效的 fd
        full_fd = result_fd.group(0)  # 提取完整 fd
        
        # 如果该 fd 还未添加为节点，则添加
        if full_fd not in G:
            G.add_node(full_fd, label=full_fd)

        # 将完整 fd 存入 fds 列表
        fd_dict[full_fd] = full_fd 
        continue
    
    first_match = None
    param_fds = param_fds_pattern.findall(str(row['Params']))
    for i, fd in enumerate(param_fds):
        if fd not in fd_dict:
            continue  # 如果 fd 不在 fd_dict 中，跳过

        # 只记录第一个有效匹配的 fd
        if first_match is None:
            first_match = fd  # 记录第一个匹配的 fd
            current_node = fd_dict[fd]  # 从叶子节点出发
                    # 创建新的节点名
            new_node = f"{fd}-{row['Function']}"
            # 如果新节点不存在，则添加它
            if new_node not in G:
                G.add_node(new_node)
            # 添加从当前节点到新节点的边（支持自环）
            G.add_edge(current_node, new_node, fd=fd)
            # 只有第一个文件描述符需要更新叶子节点的映射
            fd_dict[fd] = new_node
        else:
            current_node = fd
            G.add_edge(current_node,fd_dict[first_match])
# 绘制图
plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G)  # 使用 spring 布局

# 绘制节点和边
nx.draw_networkx_nodes(G, pos, node_size=700, node_color='lightblue', alpha=0.8)
nx.draw_networkx_edges(G, pos, edge_color='gray', arrows=True)

# 绘制节点标签
nx.draw_networkx_labels(G, pos, font_size=10, font_family="sans-serif")

# 绘制边标签（显示 fd）
edge_labels = nx.get_edge_attributes(G, 'fd')
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)

# 显示图
plt.title("FD-Based Function Call Graph")
plt.savefig('/home/chenjiaxiong/jx_code/a.png')

# 可选：保存图为 PNG 文件
G.graph['name'] = 'fd_graph'
nx.write_gexf(G, "/home/chenjiaxiong/jx_code/syscall-traffic/csv2graphfd_graph.gexf")  # 可用于Gephi等工具进一步分析

print("图构建完成，并已保存为 fd_graph.gexf 文件。")
